/*******************************************************************************

                                REGRID 

  Title:        Transform Solve output to sampled electron densities;
		Write a .map file.
  Author:       Hanna Szoke
  Date:         10/14/92
  Function:     This program takes as input a solution from Solve or Back
		(# electrons/ voxel, representing Gaussian peaks) and produces 
		a sampled electron density map in el/cub units, at a resolution 
		that is N times greater, where N is an integer (by default, 2).  
		N (spread_fac) is optionally read from the execute line:
		
			speden regrid param_file vox_file [N]

                Program expects to find a parameter file (param_file.inp) and
		a binary file named vox_file.bin. 

                If HighRes is in effect, the non-gridded input list will be 
                read from an ascii file named vox_file.list.

		For Regrid, you may specify in the .inp file that the range in 
		(x,y,z) is to span any fractional part of the unit cell, 
		including a fraction that exceeds 1.  The range (but not 
		necessarily the end points) must be positive.  

		if HDFOUT is set, Regrid writes output to group 'map' in the 
		output HDF.  otherwise Regrid writes vox_file_[N].map 

*******************************************************************************/

#include "util.h"
#include "cellparams.h"
#include "dims.h"
#include "symmetry.h"
#include "sp.h"

typedef struct {
   real	beg, end ;	/* start and end fractional coords */
} FLIMITS ;

typedef struct {
   int	beg, end ;		/* start and end integer indices */
} LIMITS ;

static	real	*hrNump ;			/* for hr points */

static	int	spread_fac = 2 ;	/* default regridding factor */
static	int     Nxin, Nyin, Nzin ;      /* original solution space dimensions */
static	int     Nxout, Nyout, Nzout ;   /* output space dimensions */
static	int     Npout ;                 /* product of Nxout*Nyout*Nzout */
static	LIMITS	rlimx, rlimy, rlimz ;	/* coord limits after regrid */
static	FLIMITS	flimx, flimy, flimz ;	/* input fractional xyz limits	*/
static	FLIMITS	flime ;              	/* input fractional data limits	*/
static	char	erange ;		/* T/F ? - apply user's range (flime) */
static	real	*dnump ;		/* regridded electron densities */
static	real	*rnump ;		/* electron per voxel array read in */
static	real	*dout ;			/* array for output  */
static	char	vox_filename[MAXSTRING] ;
static	char	list_filename[MAXSTRING] ;
real	ar_min ;
real	ar_max ;

static	void	analyzeRange() ;	/* Find range for output density map */
static	void	spread() ;		/* Distribute input over larger grid */
static	void	redefine_dimensions() ;	/* Internal and output dimensions */
static	void	r_readInput() ;		/* Read input, set up run parameters */
static	void	hrRegInput() ;		/* High-res regrid input	*/
static	void	prepareOut() ;		/* Transfer data to output array */
static	void	combine_octants(COMPLEX *, real) ;
static	void	r_echo() ;
static	void	setupXmap() ;

void	transform(int, real *, real *) ;		
					/* Do the actual regridding (FFT's) */

void	regrid_main(int argc, char *argv[])
{
     void	strip_suffix(char *, char *, char *) ;
     void	readEpixFile(char *, real *) ; 
     void	write_HDFmap(char *, real *, real, real) ;
     void       writeEpixFile(char *, real *) ;

     char	strip_name[MAXSTRING] ;
     
     if (argc > optind+4) 	/* too many arguments */
	  ballpark(caller) ; 
     else if (argc > optind+2) {
          sprintf(vox_filename, "%s", argv[optind+2]) ;
          if (argc == optind+4) 
	       spread_fac = atoi(argv[optind+3]) ;
     } 
     else {
	  prompt("\nPlease enter filename & optionally, spread factor: ") ;
          sscanf(terminp, "%s %s", vox_filename, message) ;
	  if ((int)strlen(message) > 0)
	       spread_fac = atoi(message) ;
	  sprintf(message, "Running %s on file %s with spread factor %d", 
		  caller, vox_filename, spread_fac) ;
	  printTwice(message) ;
     }

     hello(caller) ;

  /************************************************************
  Fetch input conditions, get input data
  ************************************************************/

     r_readInput() ;

     rnump = (real *) e_malloc(Np*sizeof(real), caller) ;
     printTwice("Reading solution map ...") ;

     readEpixFile(vox_filename, rnump) ;

     if (HighRes) 
	  hrRegInput(vox_filename) ;

  /************************************************************
  Spread the gridded data over a (potentially) finer grid.         
  Transform (FFT followed by * by exp. factors, then FFT-1)       
  ************************************************************/

     redefine_dimensions() ;

     printTwice("Regridding ... ") ;

     spread() ;
     transform(Np, dnump, hrNump) ;
     prepareOut(dnump) ; 

  /************************************************************
  Output will be written in pwd rather than in directory from
  which the Solve/Back output was taken.
  ************************************************************/

     strip_path(vox_filename, pwd_name) ;
     strip_suffix(pwd_name, BINSUF, strip_name) ;
     analyzeRange() ;

  /************************************************************
  Prepare filename for electron density map (output).   
  Name includes _N where N is the spread factor.
  ************************************************************/

     sprintf(out_filename, "%s_%1d", strip_name, spread_fac) ;
     if (HDFOUT) {
          printTwice("Writing regridded solution file in HDF format.\n") ;
          write_HDFmap(out_filename, dout, ar_min, ar_max) ;
     }
     else {
          printTwice("Writing regridded solution file in XPLOR format.\n") ;
          setupXmap() ;
     }
     strcat(out_filename,"_map") ;
     writeEpixFile(out_filename, dout) ;
}

        /************************************************************
  	Transfer data from full regridded array, dnump, to dout, 
	(or from rnump if no regridding was actually done)
	using limits as defined by rlimx, rlimy and rlimz.  
	These limits are not necessarily a proper subset of the full 
	regridded array and may span the edges of the unit cell.   
        ************************************************************/

void	prepareOut(real *in) 
{
     int 	i, j, k ;
     int	n, ii, jj, kk ;
     LIMITS	ilims, jlims, klims ;

     dout = (real *) e_malloc(Npout*sizeof(real), caller) ;

     ilims = rlimx ;
     jlims = rlimy ;
     klims = rlimz ;

     while (ilims.beg < 0) {
	  ilims.beg += Nx ;
	  ilims.end += Nx ;
     }
     while (jlims.beg < 0) {
	  jlims.beg += Ny ;
	  jlims.end += Ny ;
     }
     while (klims.beg < 0) {
	  klims.beg += Nz ;
	  klims.end += Nz ;
     }

     for (n = 0, k = klims.beg; k < klims.end; k++) {
	kk = k % Nz ;
        for (j = jlims.beg; j < jlims.end; j++){
	   jj = j % Ny ;
           for (i = ilims.beg; i < ilims.end; i++, n++) {
		ii = i % Nx ;
                *(dout+n) = *(in + ii + jj*Nx + kk*Nx*Ny) ;
	   } 
        }
     }
}

 	/************************************************************
	Spread elements of rnump onto a finer grid, dnump, which will
	contain sampled electrons per cubic units.  Spread uses the
	given spread_fac.  
	************************************************************/

void	spread()	
{
     int 	i, j, k ;
     int 	ii, jj, kk ;
     real 	*prnump ;
     real	*pdnump ;

     dnump = (real *) e_malloc(Np*sizeof(real), caller) ;

     prnump = rnump ;
     pdnump = dnump ;

     for (k = 0; k < Np; k++)
	  *(pdnump++) = 0 ;

     for (k = 0; k < Nzin; k++) {
	kk = spread_fac*k ;
        for (j = 0; j < Nyin; j++){
	   jj = spread_fac*j ;
           for (i = 0; i < Nxin; i++) {
	        ii = spread_fac * i ; 
                *(dnump + ii + jj*Nx + kk*Nx*Ny) = *(prnump++) ;
	   } 
        }
     }

     e_free(Np*sizeof(float), rnump) ;
}

	/************************************************************
	Transform, apply exponential Gaussian factors, and then 
	back-transform densities.  
	In contrast to the regular Solve transformations to/from 
	Fourier space, there is no need to expand the COMPLEX R to Nfob.
	Thus, if there is high resolution, we redefine Nh, Nk and Nl 
	to be Nx, Ny and Nz, respectively so that we can use doSFT to
	accumulate the list contributions into the R array.
	************************************************************/

void transform(int N, real *dnump, real *hrNump)          
{
     void	reg_initfft() ;
     void	reg_do_fft(COMPLEX *, int) ;
     void    	hrResetIndices(int);
     void	hrDoregSFT(COMPLEX *, real *);	

     COMPLEX 	*R, *Rhigh = NULL ;
     COMPLEX	*pR, *pRhigh ;
     real	*pdnump ;
     real	normfac ;
     int        n ;

     R = (COMPLEX *) e_malloc(Np*sizeof(COMPLEX), caller) ;

     /* Feed the input into R */

     for (n=0, pdnump=dnump, pR=R; n < Np; n++, pR++, pdnump++) {
          pR->re = *pdnump ;
          pR->im = 0 ;
     }

     reg_initfft() ;

     reg_do_fft(R, -1) ;			/* FORWARD FFT */

     if (HighRes) {
          hrResetIndices(spread_fac);

          Rhigh = (COMPLEX *) e_malloc(Np*sizeof(COMPLEX), caller) ;

          for (n=0, pRhigh=Rhigh; n < Np; n++, pRhigh++) {
               pRhigh->re = 0 ;
               pRhigh->im = 0 ;
          }

          hrDoregSFT(Rhigh, hrNump) ;
     }

     /* Using data from full ellipsoid, apply the spread factor */

     combine_octants(R, delfac) ;

     if (HighRes) {
          combine_octants(Rhigh, delfac/4.) ;

          for (n=0, pR = R, pRhigh=Rhigh; n < Np; n++, pR++, pRhigh++) {
               pR->re += pRhigh->re ;
               pR->im += pRhigh->im ;
          }
     }

     reg_do_fft(R, 1) ;		/* BACK FFT */

  /************************************************************
   Retrieve the real part and renormalize as indicated. 
  ************************************************************/

     normfac = crys_vol / Nptotal ;

     for (n=0, pdnump=dnump, pR=R; n < Np; n++, pR++, pdnump++) 
          *pdnump = pR->re * normfac ;

     e_free(Np*sizeof(COMPLEX), R) ; 
     e_free(Np*sizeof(COMPLEX), Rhigh) ; 
}

	/***********************************************************
	Bring together the exponential function contributions from 
	the 8 octants of the full ellipsoid.  This differs from 
	sum_over_octants, in that R is NOT expanded to the half-
	ellipsoid here;  instead, the factors that would apply if
	it were so expanded are used to "combine" the data.
	***********************************************************/

void	combine_octants(COMPLEX *R, real dfac)
{
     real     one_expfac ;
     int	n, h, k, l ;

     for (n = 0; n < Np; n++) {

          h = n % Nx ;
	  k = (n/Nx) % Ny ;
	  l = n / (Nx*Ny) ;

          one_expfac = exp(-(dfac * nusq(h, k, l))) ;

	  if (k > 0) 
               one_expfac += exp(-(dfac * nusq(h, k-Ny, l))) ;
       
	  if (l > 0)
               one_expfac += exp(-(dfac * nusq(h, k, l-Nz))) ;
           
	  if ((k > 0) && (l > 0)) 
               one_expfac += exp(-(dfac * nusq(h, k-Ny, l-Nz))) ;
           
	  if (h > 0) 
               one_expfac += exp(-(dfac * nusq(h-Nx, k, l))) ;
        
	  if ((h > 0) && (k > 0)) 
               one_expfac += exp(-(dfac * nusq(h-Nx, k-Ny, l))) ;
         
	  if ((h > 0) && (l > 0))
               one_expfac += exp(-(dfac * nusq(h-Nx, k, l-Nz))) ;
          
	  if ((h > 0) && (k > 0) && (l > 0)) 
               one_expfac += exp(-(dfac * nusq(h-Nx, k-Ny, l-Nz))) ;
	   
          (R+n)->re *= one_expfac ;
          (R+n)->im *= one_expfac ;

      }
}
	 
void r_readInput()    /*   Read Regrid input	*/ 
{
     void	getRange(int, real *, real *, real *, real *, 
			 real *, real *) ;
     int	k ;
     int	status = 0 ;

     readBasicInput() ;

     Nxin = Nx ;
     Nyin = Ny ;
     Nzin = Nz ;

     /* Defaults for optional and some required input.   */

     flimx.beg = 0 ; flimx.end = 1 ;
     flimy.beg = 0 ; flimy.end = 1 ;
     flimz.beg = 0 ; flimz.end = 1 ;
     flime.beg = 0 ; flime.end = 1 ;
     erange = FALSE ;

     for (k = 0; k < Nlines; k++) {
          strcpy(nextline, allinp + k*MAXSTRING) ;
          sscanf(nextline, "%s", id) ;

          if (strcmp(id, "X_LIMITS") == 0) {
               sscanf(nextline, "%*s %g %g", &t1, &t2) ;
	       flimx.beg = t1 ; 
	       flimx.end = t2 ;
          }
          else if (strcmp(id, "Y_LIMITS") == 0) {
               sscanf(nextline, "%*s %g %g", &t1, &t2) ;
	       flimy.beg = t1 ; 
	       flimy.end = t2 ;
          }
          else if (strcmp(id, "Z_LIMITS") == 0) {
               sscanf(nextline, "%*s %g %g", &t1, &t2) ;
	       flimz.beg = t1 ; 
	       flimz.end = t2 ;
          }
          else if (strcmp(id, "E_LIMITS") == 0) {
               sscanf(nextline, "%*s %g %g", &t1, &t2) ;
	       flime.beg = t1 ; 
	       flime.end = t2 ;
          }
     }

     /************************************************************
     Be sure spread_fac is reasonable & consistent with grid type 
     ************************************************************/

     if (spread_fac > 8)
          Fatal("Your spread factor is too large!") ;

     /************************************************************
     Check for legality of dimensional input
     ************************************************************/

     if ((flimx.end - flimx.beg) < 0)  {
	 Warning("Negative x_limits span in input.") ;
         status = -1 ;
     }
     if ((flimy.end - flimy.beg) < 0)  {
	 Warning("Negative y_limits span in input.") ;
         status = -1 ;
     }
     if ((flimz.end - flimz.beg) < 0)  {
	 Warning("Negative z_limits span in input.") ;
         status = -1 ;
     }
     if ((flime.end - flime.beg) < 0)  {
	 Warning("Negative electron limits span in input.") ;
         status = -1 ;
     }

     if (status == -1)
	  Fatal("Quitting.") ;

     if ((flime.beg != 0) || (flime.end != 1))
	  erange = TRUE ;

     r_echo() ;
}
	 
void	hrRegInput(char *vox_filename)
{
     int	hrGetLength(char *) ;
     void	hrReadLists(char *, real *) ;

     sprintf(list_filename, "%s.list", vox_filename) ;
     Nhr = hrGetLength(list_filename) ;
     hrNump = (real *) e_malloc(Nhr*sizeof(real), caller) ;
     hrReadLists(list_filename, hrNump) ;	
     sprintf(message, "including %d high-resolution points.", Nhr) ;
     printTwice(message) ;
}


void redefine_dimensions()
{
     /* Redefine internal (regridded) dimensions */

     Nx = spread_fac * Nxin ;
     Ny = spread_fac * Nyin ;
     Nz = spread_fac * Nzin ;
     Np = Nx * Ny * Nz ;

     /* Translate user's limits into limits in the output array */

     rlimx.beg = (floor) (Nx * flimx.beg) ;
     rlimy.beg = (floor) (Ny * flimy.beg) ;
     rlimz.beg = (floor) (Nz * flimz.beg) ;

     rlimx.end = (ceil) (Nx * flimx.end) ;
     rlimy.end = (ceil) (Ny * flimy.end) ;
     rlimz.end = (ceil) (Nz * flimz.end) ;

     Nxout = (rlimx.end - rlimx.beg) ;
     Nyout = (rlimy.end - rlimy.beg) ;
     Nzout = (rlimz.end - rlimz.beg) ;

     Npout = Nxout * Nyout * Nzout ;

     if ((flimx.beg != 0) || (flimx.end != 1) ||
         (flimy.beg != 0) || (flimy.end != 1) ||
         (flimz.beg != 0) || (flimz.end != 1)) {

          sprintf(message,
	    "\nThe actual gridded coordinate ranges in %s follow -", units) ;
          printTwice(message) ;
          sprintf(message,
            "\tx limits: (%g, %g)", aAxis*rlimx.beg / Nx, aAxis*rlimx.end / Nx);
          printTwice(message) ;
          sprintf(message,
            "\ty limits: (%g, %g)", bAxis*rlimy.beg / Ny, bAxis*rlimy.end / Ny);
          printTwice(message) ;
          sprintf(message,
            "\tz limits: (%g, %g)", cAxis*rlimz.beg / Nz, cAxis*rlimz.end / Nz);
          printTwice(message) ;

          printTwice("\n") ;
     }
}

void r_echo()      /*  Report input values to output. */ 
{

     printTwice("\n") ;

     sprintf(message, 
          "Unit cell measures %6.2f by %6.2f by %6.2f %s", 
	  aAxis, bAxis, cAxis, units);
     printTwice(message) ;
     sprintf(message, "Eta, exp. factor, is %g", eta) ;
     printTwice(message) ;
     sprintf(message, "Regridding factor (spread_fac) is %d", spread_fac) ;
     printTwice(message) ;

     sprintf(message, "\nThe requested coordinate ranges in %s follow -", units) ;
     printTwice(message) ;
     sprintf(message, 
	       "\tx limits: (%g, %g)", aAxis*flimx.beg, aAxis*flimx.end) ;
     printTwice(message) ;
     sprintf(message, 
	       "\ty limits: (%g, %g)", bAxis*flimy.beg, bAxis*flimy.end) ;
     printTwice(message) ;
     sprintf(message, 
	       "\tz limits: (%g, %g)", cAxis*flimz.beg, cAxis*flimz.end) ;
     printTwice(message) ;
     if (erange)
          sprintf(message, 
		"\tRange of output will be %g - %g.", flime.beg, flime.end);
     else
          sprintf(message,
		"\tRange of output will be in units of el/cubic %s.", units) ;

     printTwice(message) ;
}

#include	<sys/time.h>	/* ... for picking up date & time	*/

	/************************************************************
	Write a file of output electron densities in XPLOR format.
	************************************************************/

void	setupXmap()	
{
     void	printXmap();		/* Write output in XPLOR format */
     char	Range[MAXSTRING]; 

     sprintf(Range," %7d %7d %7d %7d %7d %7d %7d %7d %7d",
	     Nx, rlimx.beg, rlimx.end-1,
	     Ny, rlimy.beg, rlimy.end-1, 
	     Nz, rlimz.beg, rlimz.end-1) ;

     printXmap(out_filename, Range, Nxout, Nyout, Nzout, dout) ;
}

void	analyzeRange()	/* Analyze a solution map  */
{
     int 	get_min_max(real *, int, real *, real *) ;
     real	scale_factor ;	
     int	n ;

     get_min_max(dout, Npout, &ar_min, &ar_max) ;

     if (erange) {
          if (ar_min == ar_max)
               Fatal("Output range is zero!") ;
          else {
               scale_factor = (real)(flime.end - flime.beg) / (ar_max - ar_min) ;

               for (n = 0; n < Npout; n++) 
	            *(dout+n) = flime.beg + scale_factor * (*(dout+n) - ar_min) ;
          }
     }
     else {
          sprintf(message,
             "\nRange of output is (%g, %g) el/cub %s", ar_min, ar_max, units) ;
          printTwice(message) ;
     }
     return ;
}
